module io

    use global
    use netcdf

    implicit none
    private

    public :: init, pos, lat, lon, depth 

    real, allocatable, dimension(:) :: lat, lon, depth

    contains

    subroutine init()

        integer :: ncid, status
        integer :: u_id, v_id, w_id, lat_id, lon_id, depth_id
        character(len=1) :: U_NAME = "U"
        character(len=1) :: V_NAME = "V"
        character(len=2) :: W_NAME = "WT"
        character(len=8) :: LON_NAME = "XU_OCEAN"
        character(len=8) :: LAT_NAME = "YU_OCEAN"
        character(len=8) :: DEPTH_NAME = "ZW_OCEAN"

        ! This is for B-grid && MOM4.
        real, allocatable, dimension(:,:,:,:) :: uvel, vvel, wvel, u, v, w
        real, allocatable, dimension(:) :: lon2, lat2, depth2

        real :: u_missing, v_missing, w_missing

        integer :: i, j, k, l, lc, lf
        real :: dlon, dlat, dH, cx, cy, ds
        real :: F(2,2)

        namelist /config_nml/ idm, jdm, kdm, ldm, filename, transport, direction, maxsteps, interval, cyclic_boundaries

        ! Read configuration from namelist.
        open(1, file="config.nml")
        read(1, nml=config_nml)
        close(1)

        allocate(ax(idm, jdm, kdm, ldm))
        allocate(bx(idm, jdm, kdm, ldm))
        allocate(gx(idm, jdm, kdm, ldm))
        allocate(dx(idm, jdm, kdm, ldm))
        allocate(ay(idm, jdm, kdm, ldm))
        allocate(by(idm, jdm, kdm, ldm))
        allocate(gy(idm, jdm, kdm, ldm))
        allocate(dy(idm, jdm, kdm, ldm))
        allocate(az(idm, jdm, kdm, ldm))
        allocate(bz(idm, jdm, kdm, ldm))
        allocate(gz(idm, jdm, kdm, ldm))
        allocate(dz(idm, jdm, kdm, ldm))
        allocate(dv(idm, jdm, kdm, ldm))
        allocate(lon(idm))
        allocate(lat(jdm))
        allocate(depth(kdm))
        allocate(volume(idm, jdm, kdm, 1))

        allocate(uvel(idm, jdm, kdm, ldm))
        allocate(vvel(idm, jdm, kdm, ldm))
        allocate(wvel(idm, jdm, kdm, ldm))
        allocate(u(0:idm, 0:jdm, kdm, ldm))
        allocate(v(0:idm, 0:jdm, kdm, ldm))
        allocate(w(idm, jdm, 0:kdm, ldm))
        allocate(lon2(0:idm))
        allocate(lat2(0:jdm))
        allocate(depth2(0:kdm))

        ! Set number of time values for the depth variable. For MOM4 this is 1,
        ! since depth is constant. For MICOM, this should be equal to the number
        ! of time steps in the data (ldm).
        dldm = 1

        ! Open file.
        status = nf90_open(filename, nf90_nowrite, ncid)

        ! Read grid.
        status = nf90_inq_varid(ncid, LON_NAME, lon_id)
        status = nf90_get_var(ncid, lon_id, lon)
        status = nf90_inq_varid(ncid, LAT_NAME, lat_id)
        status = nf90_get_var(ncid, lat_id, lat)
        status = nf90_inq_varid(ncid, DEPTH_NAME, depth_id)
        status = nf90_get_var(ncid, depth_id, depth)

        ! Read velocities.
        status = nf90_inq_varid(ncid, U_NAME, u_id)
        status = nf90_get_var(ncid, u_id, uvel)
        status = nf90_inq_varid(ncid, V_NAME, v_id)
        status = nf90_get_var(ncid, v_id, vvel)
        status = nf90_inq_varid(ncid, W_NAME, w_id)
        status = nf90_get_var(ncid, w_id, wvel)

        ! Read missing values for the velocities.
        status = nf90_get_att(ncid, u_id, "missing_value", u_missing)
        status = nf90_get_att(ncid, v_id, "missing_value", v_missing)
        status = nf90_get_att(ncid, w_id, "missing_value", w_missing)

        ! Close file.
        status = nf90_close(ncid)

        ! Replace missing values with zero.
        do i = 1, idm
        do j = 1, jdm
        do k = 1, kdm
        do l = 1, ldm
            if (uvel(i,j,k,l) == u_missing) uvel(i,j,k,l) = 0.0
            if (vvel(i,j,k,l) == v_missing) vvel(i,j,k,l) = 0.0
            if (wvel(i,j,k,l) == w_missing) wvel(i,j,k,l) = 0.0
        enddo
        enddo
        enddo
        enddo

        ! Calculate u, v and w from velocities. These will be used to calculate
        ! the fluxes.
        u(1:idm,1:jdm,:,:) = uvel(:,:,:,:) * direction
        v(1:idm,1:jdm,:,:) = vvel(:,:,:,:) * direction
        w(:,:,1:kdm,:) = -wvel(:,:,:,:) * direction  ! invert to coincide with k axis

        ! Pad velocities with interpolated values.
        u(0,1:jdm,:,:) = 2*u(1,1:jdm,:,:) - u(2,1:jdm,:,:)
        u(:,0,:,:) = 2*u(:,1,:,:) - u(:,2,:,:)
        v(0,1:jdm,:,:) = 2*v(1,1:jdm,:,:) - v(2,1:jdm,:,:)
        v(:,0,:,:) = 2*v(:,1,:,:) - v(:,2,:,:)
        w(:,:,0,:) = 0.0                 ! surface

        ! Pad lat/lon/depth with interpolated values.
        lon2(1:idm) = lon(:)
        lon2(0) = 2*lon(1) - lon(2)
        lat2(1:jdm) = lat(:)
        lat2(0) = 2*lat(1) - lat(2)
        depth2(1:kdm) = depth(:)
        depth2(0) = 0.0

        ! Calculate coefficients from velocities.
        ! We first need to calculate the length of the box: dlon * dlat * dH
        cy = R * PI/180.0
        do j = 1,jdm
            cx = cy * cos(0.5*(lat2(j)+lat2(j-1))*pi/180.0)
            dlat = cy * abs(lat2(j) - lat2(j-1))
            do i = 1,idm
                dlon = cx * abs(lon2(i) - lon2(i-1))
                do k = 1,kdm
                    ! Calculate layer thickness.
                    dH = depth2(k) - depth2(k-1)

                    ! Scaled time ds = dt/volume
                    volume(i,j,k,1) = dH * dlon * dlat
                    ds = interval/volume(i,j,k,1)

                    ! Get u velocity at cell corners [u(i,j,k,l), nw, sw, se].
                    do lc = 1,ldm
                        ! Index for future time step (lf), as opposed to current (lc).
                        ! 1 -> 2, 4 -> 5, ldm -> 1
                        lf = 1 + mod(lc, ldm)

                        ! Coefficients in x.
                        F(1,1) = 0.5*(u(i-1,j,k,lc) + u(i-1,j-1,k,lc)) * dlat * dH
                        F(1,2) = 0.5*(u(i-1,j,k,lf) + u(i-1,j-1,k,lf)) * dlat * dH
                        F(2,1) = 0.5*(u(i,j,k,lc) + u(i,j-1,k,lc)) * dlat * dH
                        F(2,2) = 0.5*(u(i,j,k,lf) + u(i,j-1,k,lf)) * dlat * dH
                        ax(i,j,k,lc) = -(F(2,2) - F(1,2) - F(2,1) + F(1,1))/ds
                        bx(i,j,k,lc) = F(1,1) - F(2,1)
                        gx(i,j,k,lc) = -(F(1,2) - F(1,1))/ds
                        dx(i,j,k,lc) = -F(1,1)

                        ! Coefficients in y.
                        F(1,1) = 0.5*(v(i,j-1,k,lc) + v(i-1,j-1,k,lc)) * dlon * dH
                        F(1,2) = 0.5*(v(i,j-1,k,lf) + v(i-1,j-1,k,lf)) * dlon * dH
                        F(2,1) = 0.5*(v(i,j,k,lc) + v(i-1,j,k,lc)) * dlon * dH
                        F(2,2) = 0.5*(v(i,j,k,lf) + v(i-1,j,k,lf)) * dlon * dH
                        ay(i,j,k,lc) = -(F(2,2) - F(1,2) - F(2,1) + F(1,1))/ds
                        by(i,j,k,lc) = F(1,1) - F(2,1)
                        gy(i,j,k,lc) = -(F(1,2) - F(1,1))/ds
                        dy(i,j,k,lc) = -F(1,1)

                        ! And in z.
                        F(1,1) = w(i,j,k-1,lc) * dlon * dlat
                        F(1,2) = w(i,j,k-1,lf) * dlon * dlat
                        F(2,1) = w(i,j,k,lc) * dlon * dlat
                        F(2,2) = w(i,j,k,lf) * dlon * dlat
                        az(i,j,k,lc) = -(F(2,2) - F(1,2) - F(2,1) + F(1,1))/ds
                        bz(i,j,k,lc) = F(1,1) - F(2,1)
                        gz(i,j,k,lc) = -(F(1,2) - F(1,1))/ds
                        dz(i,j,k,lc) = -F(1,1)
                    enddo
                enddo
            enddo
        enddo

    end subroutine init

    subroutine pos(i, j, k, l, rx, ry, rz, x, y, z)

        integer, intent(in) :: i, j, k, l
        real, intent(in) :: rx, ry, rz
        real, intent(out) :: x, y, z

        if (i > 1) then
            x = lon(i-1) + rx*abs(lon(i) - lon(i-1))
        else
            x = lon(1) - (1-rx)*abs(lon(2) - lon(1))
        endif
        if (j > 1) then
            y = lat(j-1) + ry*abs(lat(j) - lat(j-1))
        else
            y = lat(j) - (1-ry)*abs(lat(2) - lat(1))
        endif
        if (k > 1) then
            z = depth(k-1) + rz*abs(depth(k) - depth(k-1))
        else
            z = rz*depth(1)
        endif

    end subroutine pos

end module io
